﻿#ifdef WIN32

#include "stack_trace_windows.hh"
#include "../box/service_box.hh"
#include "../box/system_exception.hh"
#include "../util/string_util.hh"
#include <iomanip>
#include <memory>
#include <signal.h>
#include <sstream>
#include <string>
#include <vector>
#include <windows.h>
#include <dbghelp.h>
#include <eh.h>

namespace kratos {
namespace util {

bool global_open_system_exception = false;

auto get_stack_frame_info_internal(
  PSYMBOL_INFO symbol,
  DWORD64 address,
  std::ostream &ost,
  HANDLE process,
  int depth) {
  IMAGEHLP_LINE64 image_line64;
  image_line64.SizeOfStruct = sizeof(IMAGEHLP_LINE64);
  DWORD64 sym = 0;
  DWORD _na = 0;
  ost << std::left << std::setw(2) << depth;
  // 打印每一帧的调用函数，文件名，行号
  if (::SymFromAddr(process, address, &sym, symbol) &&
      ::SymGetLineFromAddr64(process, address, &_na, &image_line64)) {
    if (image_line64.FileName) {
      // 帧深度，地址，函数名，文件名，行号
      ost << "  0x" << std::hex << symbol->Address << std::dec << " "
          << symbol->Name << " (" << image_line64.FileName << ":"
          << image_line64.LineNumber << ")" << std::endl;
    }
  } else {
    // 无法获取，无符号表
    ost << "  0x? ?" << std::endl;
  }
}

auto reset_symbol(char *symbol_buffer, int size) -> PSYMBOL_INFO {
  std::memset(symbol_buffer, 0, size);
  auto *symbol = (PSYMBOL_INFO)symbol_buffer;
  symbol->SizeOfStruct = sizeof(SYMBOL_INFO);
  symbol->MaxNameLen = MAX_SYM_NAME;
  return symbol;
}

auto reset_symbol64(char *symbol_buffer, int size) -> IMAGEHLP_SYMBOL64 * {
  std::memset(symbol_buffer, 0, size);
  auto *image_symbol = (IMAGEHLP_SYMBOL64 *)symbol_buffer;
  image_symbol->SizeOfStruct = sizeof(IMAGEHLP_SYMBOL64);
  image_symbol->MaxNameLength = MAX_SYM_NAME;
  return image_symbol;
}

auto StackTrace::get_stack_frame_info(int max_frames) -> std::string {
  constexpr static int BUFFER_SIZE =
      sizeof(SYMBOL_INFO) + MAX_SYM_NAME * sizeof(TCHAR);
  std::unique_ptr<char> buffer(new char[BUFFER_SIZE]);
  // 当前进程句柄
  auto process = ::GetCurrentProcess();
  // 初始化
  ::SymInitialize(process, nullptr, TRUE);
  ::SymSetOptions(SYMOPT_LOAD_LINES);
  std::unique_ptr<void *> stack_frames(new void *[max_frames]);
  // 获取帧地址
  auto frames =
    ::CaptureStackBackTrace(0, max_frames, stack_frames.get(), nullptr);
  std::stringstream ss;
  ss << std::endl;
  // 帧深度
  auto depth = (int)frames - 1;
  for (int i = 0; i < (int)frames; i++) {
    auto symbol = reset_symbol(buffer.get(), BUFFER_SIZE);
    get_stack_frame_info_internal(
      symbol,
      (DWORD64)(stack_frames.get()[i]),
      ss,
      process,
      depth--
    );
  }
  ::SymCleanup(process);
  return ss.str();
}

auto get_stack_trace_info_from_exception_record(CONTEXT *current_context)
    -> std::string {
  constexpr static int BUFFER_SIZE =
    sizeof(IMAGEHLP_SYMBOL64) + MAX_SYM_NAME * sizeof(TCHAR);
  char seh_buffer[BUFFER_SIZE] = {0};
  STACKFRAME64 stack_frame64;
  std::vector<std::string> frame_vector;
  std::stringstream output_stream;
  // 当前进程句柄
  auto process = ::GetCurrentProcess();
  auto thread = ::GetCurrentThread();
  // 初始化
  ::SymInitialize(process, nullptr, TRUE);
  ::SymSetOptions(SYMOPT_LOAD_LINES);
  // 建立一个新的，防止在StackWalk64内部修改
  CONTEXT context;
  memcpy(&context, current_context, sizeof(CONTEXT));
#if _M_X64 // TODO 现在只有x64配置
  DWORD dwImageType = IMAGE_FILE_MACHINE_AMD64;
  stack_frame64.AddrPC.Offset = context.Rip;
  stack_frame64.AddrPC.Mode = AddrModeFlat;
  stack_frame64.AddrFrame.Offset = context.Rsp;
  stack_frame64.AddrFrame.Mode = AddrModeFlat;
  stack_frame64.AddrStack.Offset = context.Rsp;
  stack_frame64.AddrStack.Mode = AddrModeFlat;
#else
#error "Platform undefined"
#endif // _M_X64
  std::memset(&stack_frame64, 0, sizeof(STACKFRAME64));
  IMAGEHLP_LINE64 image_line64;
  image_line64.SizeOfStruct = sizeof(IMAGEHLP_LINE64);
  for (;;) {
    output_stream.str("");
    if (!::StackWalk64(
      dwImageType,
      process,
      thread,
      &stack_frame64,
      &context,
      nullptr,
      SymFunctionTableAccess64,
      SymGetModuleBase64,
      nullptr)) {
      break;
    }
    if (stack_frame64.AddrFrame.Offset == 0) {
      break;
    }
    auto symbol = reset_symbol64(seh_buffer, BUFFER_SIZE);
    DWORD64 sym = 0;
    DWORD _na = 0;
    if (::SymGetSymFromAddr64(
      process,
      stack_frame64.AddrPC.Offset,
      nullptr,
      symbol) &&
        ::SymGetLineFromAddr64(
          process,
          stack_frame64.AddrPC.Offset,
          &_na,
          &image_line64)) {
      if (image_line64.FileName) {
        // 帧深度，地址，函数名，文件名，行号
        output_stream << "  0x" << std::hex << symbol->Address << std::dec
                      << " " << util::demangle(symbol->Name) << " ("
                      << image_line64.FileName << ":" << image_line64.LineNumber
                      << ")" << std::endl;
        frame_vector.emplace_back(output_stream.str());
      }
    } else {
      // 无法获取，无符号表
      output_stream << "  0x? ?" << std::endl;
      frame_vector.emplace_back(output_stream.str());
    }
  }
  ::SymCleanup(process);
  output_stream.str("");
  output_stream << std::endl;
  // 按照堆栈深度记录信息
  auto depth = frame_vector.size() - 1;
  for (const auto &s : frame_vector) {
    output_stream << std::left << std::setw(2) << depth-- << s;
  }
  return output_stream.str();
}

auto StackTrace::get_unhandled_exception_stack_frame_info(void *except_info,
                                                          int max_frames)
    -> std::string {
  EXCEPTION_POINTERS *ptr = (EXCEPTION_POINTERS *)except_info;
  // 获取异常发生时的堆栈
  return get_stack_trace_info_from_exception_record(ptr->ContextRecord);
}

kratos::service::ServiceBox *global_box_ = nullptr;
LPTOP_LEVEL_EXCEPTION_FILTER default_seh_handler = nullptr;
_se_translator_function default_seh_translator = nullptr;

auto seh_handler_internal(EXCEPTION_POINTERS *exception_ptr) -> LONG {
  auto stack_trace_info =
      get_stack_trace_info_from_exception_record(exception_ptr->ContextRecord);
  if (global_box_) {
    global_box_->write_log_line(klogger::Logger::FATAL, stack_trace_info);
  } else {
    std::cerr << stack_trace_info << std::endl;
  }
  return EXCEPTION_EXECUTE_HANDLER;
}

auto seh_exception_translator(unsigned int exception_code,
                              EXCEPTION_POINTERS *exception) -> void {
  auto stack_info =
      get_stack_trace_info_from_exception_record(exception->ContextRecord);
  switch (exception_code) {
  case EXCEPTION_ACCESS_VIOLATION:
    throw service::SegmentationFaultException(stack_info.c_str());
    break;
  case EXCEPTION_FLT_DIVIDE_BY_ZERO:
  case EXCEPTION_INT_DIVIDE_BY_ZERO:
  case EXCEPTION_FLT_DENORMAL_OPERAND:
  case EXCEPTION_FLT_INEXACT_RESULT:
  case EXCEPTION_FLT_INVALID_OPERATION:
  case EXCEPTION_FLT_OVERFLOW:
  case EXCEPTION_FLT_STACK_CHECK:
  case EXCEPTION_FLT_UNDERFLOW:
  case EXCEPTION_INT_OVERFLOW:
    throw service::NumberException(stack_info.c_str());
    break;
  case EXCEPTION_STACK_OVERFLOW:
    throw service::StackFaultException(stack_info.c_str());
    break;
  case EXCEPTION_ARRAY_BOUNDS_EXCEEDED:
    throw service::ArrayException(stack_info.c_str());
    break;
  case EXCEPTION_BREAKPOINT:
  case EXCEPTION_SINGLE_STEP:
    throw service::DebugException(stack_info.c_str());
    break;
  default:
    break;
  }
  throw service::SystemException(stack_info.c_str());
}

void abort_handler(int) {
  auto stack_trace_info = StackTrace::get_stack_frame_info();
  if (global_open_system_exception) {
    throw service::AbortException(stack_trace_info.c_str());
  } else {
    if (global_box_) {
      global_box_->write_log_line(klogger::Logger::FATAL, stack_trace_info);
    } else {
      std::cerr << stack_trace_info << std::endl;
    }
  }
}

_crt_signal_t abrt_handler = nullptr;

auto StackTrace::install_seh_handler(kratos::service::ServiceBox *box) -> bool {
  if (global_box_) {
    return false;
  }
  global_box_ = box;
  // 设置自定义顶层处理器
  default_seh_handler =
      ::SetUnhandledExceptionFilter((LPTOP_LEVEL_EXCEPTION_FILTER)seh_handler);
  if (!abrt_handler) {
    // 设置abort处理器
    ::_set_abort_behavior(0, _WRITE_ABORT_MSG);
    ::_set_abort_behavior(0, _CALL_REPORTFAULT);
    abrt_handler = ::signal(SIGABRT, abort_handler);
  }
  return true;
}

auto StackTrace::uninstall_seh_handler() -> void {
  // 恢复默认的处理器
  if (default_seh_handler) {
    ::SetUnhandledExceptionFilter(default_seh_handler);
  }
  if (abrt_handler) {
    ::signal(SIGABRT, abrt_handler);
  }
}

auto StackTrace::install_system_exception(kratos::service::ServiceBox *box)
    -> bool {
  if (global_box_) {
    return false;
  }
  global_box_ = box;
  default_seh_translator = ::_set_se_translator(seh_exception_translator);
  if (!abrt_handler) {
    // 设置abort处理器
    ::_set_abort_behavior(0, _WRITE_ABORT_MSG);
    ::_set_abort_behavior(0, _CALL_REPORTFAULT);
    abrt_handler = ::signal(SIGABRT, abort_handler);
  }
  global_open_system_exception = true;
  return true;
}

auto StackTrace::uninstall_system_exception() -> void {
  ::_set_se_translator(default_seh_translator);
  if (abrt_handler) {
    ::signal(SIGABRT, abrt_handler);
  }
  global_open_system_exception = false;
}

auto StackTrace::seh_handler(void *exception_pointers) -> long {
  return seh_handler_internal((EXCEPTION_POINTERS *)exception_pointers);
}

} // namespace util
} // namespace kratos

#endif // WIN32
